import sys
import os
import torch
from NN.Model_importer.Importer import FVGN
from dataset.Load_mesh import DatasetFactory
from utils import get_param
import time
from utils.get_param import get_hyperparam
from utils.Logger import Logger
from utils.losses import LpLoss
import random
import datetime

# configurate parameters
params = get_param.params(
    f"/home/xiaoli/project/3D-ShapeNet-car/src/Carshapnet/track_A_B/Logger/net GM; hs 128;/2024-07-11-21:52:48-pre-comit/states"
)
params.load_date_time = "2024-07-11-21:52:48-pre-comit"
params.load_index = "130"
params.on_gpu = 1

random.seed(int(datetime.datetime.now().timestamp()))
torch.manual_seed(int(datetime.datetime.now().timestamp()))
torch.cuda.set_per_process_memory_fraction(0.99, params.on_gpu)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# initialize Logger and load model / optimizer if according parameters were given
logger = Logger(
    get_hyperparam(params),
    datetime=params.load_date_time,
    use_csv=False,
    use_tensorboard=False,
    copy_code=False,
)

# initialize Training Dataset
datasets_factory = DatasetFactory(
    params=params,
    device=device,
    split="test",
)

# create dataset objetc
test_dataset, test_loader, test_sampler = datasets_factory.create_testset(
    batch_size=1,
    num_workers=0,
    pin_memory=False,
    persistent_workers=False,
    valid_num=50,
)

# initialize fluid model
model = FVGN(params)

fluid_model = model.to(device)
fluid_model.eval()

params.load_date_time, params.load_index = logger.load_state(
    model=fluid_model,
    optimizer=None,
    scheduler=None,
    datetime=params.load_date_time,
    index=params.load_index,
    device=device,
)
params.load_index = params.load_index
print(f"loaded: {params.load_date_time}, {params.load_index}")

params.load_index = 0 if params.load_index is None else params.load_index

lp_loss = LpLoss(size_average=True)

# 初始化用于收集整个数据集的残差的张量
epoc_loss = 0

test_loss = 0

start = time.time()
with torch.no_grad():

    epoc_val_loss = 0

    for batch_index, (
        graph_node,
        graph_edge,
        graph_cell,
    ) in enumerate(test_loader):

        (graph_node, graph_edge, graph_cell) = test_dataset.datapreprocessing(
            graph_node.cuda(),
            graph_edge.cuda(),
            graph_cell.cuda(),
            is_training=False,
        )

        pred_node_valid = fluid_model(
            graph_node=graph_node,
            graph_edge=graph_edge,
            graph_cell=graph_cell,
            is_training=False,
            params=params,
        )

        reversed_node_press = (
            (pred_node_valid - 1e-8) * test_dataset.physics_info["pressure_min_std"][1]
        ) + test_dataset.physics_info["pressure_min_std"][0]
        # reversed_node_press = pred_node_valid

        test_loss += (
            lp_loss(reversed_node_press.unsqueeze(0), graph_node.y.unsqueeze(0))
            .cpu()
            .detach()
            .item()
        )
        diff = torch.abs(pred_node_valid - graph_node.norm_y).cpu().detach().numpy()

        res_dict = {
            "node|pos": graph_node.pos.cpu().detach().numpy(),
            "node|pressure": reversed_node_press.cpu().detach().numpy(),
            "node|pressure_gt": graph_node.y.cpu().detach().numpy(),
            "cells_node": graph_node.face[0].cpu().detach().numpy(),
            "node|diff": diff,
        }

        # logger.plot(
        #     res_dict=res_dict,
        #     data_index="".join(
        #                 [
        #                     chr(ascii_code)
        #                     for ascii_code in graph_node.origin_id.cpu().tolist()
        #                 ]
        #                     ),
        #     split="test",
        # )

        logger.save_test_results(
            value=reversed_node_press.cpu().detach().squeeze().numpy(),
            num_id="".join(
                        [
                            chr(ascii_code)
                            for ascii_code in graph_node.origin_id.cpu().tolist()
                        ]
                            ),
        )

test_loss /= len(test_loader)
print(f"test loss: {test_loss}")

print(f"Generating answer completed completed in {time.time() - start:.2f} seconds")
